from langchain.retrievers import TimeWeightedVectorStoreRetriever

from afamily.llm_engine import llm,embeddings_model

from langchain_postgres.vectorstores import PGVector

LLM = llm  # Can be any LLM you want.


import math
from langchain.vectorstores.pgvector import DistanceStrategy



def relevance_score_fn(score: float) -> float:
    """Return a similarity score on a scale [0, 1]."""
    # This will differ depending on a few things:
    # - the distance / similarity metric used by the VectorStore
    # - the scale of your embeddings (OpenAI's are unit norm. Many others are not!)
    # This function converts the euclidean norm of normalized embeddings
    # (0 is most similar, sqrt(2) most dissimilar)
    # to a similarity function (0 to 1)
    return 1.0 - score / math.sqrt(2)

CONNECTION_STRING = f"postgresql+psycopg2://testuser:testpwd@localhost:5432/vectordb"


def create_new_memory_retriever():
    """Create a new vector store retriever unique to the agent."""
    # Define your embedding model

    # Initialize the vectorstore as empty
    from langchain_core.documents import Document
    embedding_size = 1024


    collection_name = "embeddings"


    vector_store = PGVector(
        embeddings=embeddings_model,
        collection_name=collection_name,
        connection=CONNECTION_STRING,
    )
    docs = [
    Document(
        page_content="there are cats in the pond",
        metadata={"id": 100, "location": "pond", "topic": "animals"},
    ),
    
    ]

    vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])
    return TimeWeightedVectorStoreRetriever(
        vectorstore=vector_store, other_score_keys=["importance"], k=15
    )
re = create_new_memory_retriever()
print(re)